import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
# 本节重点：正则化缓解过拟合

BATCH_SIZE = 30
seed = 2
rdm = np.random.RandomState(seed)
# 用随机数生成数据集
X = rdm.randn(300,2)
# 标签制作
Y_ = [int(x0*x0 + x1*x1<2) for (x0,x1) in X]
# 遍历Y，1赋值为red，其余blue
Y_c = [['red' if y else 'blue'] for y in Y_]
# shape整理， 第一个元素为-1表示随第二个参数计算得到，第二个元素表示多少列，把X
# 整理为n行2列，把Y整理为n行1列
X = np.vstack(X).reshape(-1,2)
Y_ = np.vstack(Y_).reshape(-1,1)

print (X)
print (Y_)
print (Y_c)

# 用plt.scatter画出数据集X的(x0,x1)，用各行Y_c对应的值表示颜色
plt.scatter(X[:,0], X[:,1], c = np.squeeze(Y_c))
plt.show()



# --------------------------------------------------------------------
# 定义神经网络
def get_weight(shape, regularizer): # 正则项
    w = tf.Variable(tf.random_normal(shape),dtype = tf.float32)
    tf.add_to_collection('losses',tf.contrib.layers.l2_regularizer(regularizer)(w))
    return w

def get_bias(shape):
    b = tf.Variable(tf.constant(0.01, shape = shape))
    return b

x = tf.placeholder(tf.float32, shape = (None, 2))
y_ = tf.placeholder(tf.float32, shape = (None, 1))

w1 = get_weight([2,11], 0.01)
b1 = get_bias([11])
y1 = tf.nn.relu(tf.matmul(x, w1)+b1)

w2 = get_weight([11,1], 0.01)
b2 = get_bias([1])
y = tf.matmul(y1, w2)+b2  # 这里设置输出层不过激活

# 定义损失函数
loss_mse = tf.reduce_mean(tf.square(y-y_))
loss_total = loss_mse + tf.add_n(tf.get_collection('losses')) # 加上每一个正则化的损失



# -----------------------------------------------------------------------
# 定义反向传播方法，这里不含正则化
train_step = tf.train.AdamOptimizer(0.0001).minimize(loss_mse)

with tf.Session() as sess:
    init_op = tf.global_variables_initializer()
    sess.run(init_op)
    STEPS = 40000
    for i in range(STEPS):
        start = (i*BATCH_SIZE) % 300
        end = start + BATCH_SIZE
        sess.run(train_step, feed_dict = {x:X[start:end], y_:Y_[start:end]})
        if i% 2000 == 0:
            loss_mse_v = sess.run(loss_mse, feed_dict ={x:X,y_:Y_})
            print ("在 %d 步之后，损失达到 %f" % (i, loss_mse_v))

    xx, yy = np.mgrid[-3:3:.01, -3:3:.01]
    # 将xx，yy拉直，并合并成一个2列的矩阵，得到网格坐标点集合
    grid = np.c_[xx.ravel(),yy.ravel()]
    probs = sess.run(y, feed_dict = {x:grid})
    probs = probs.reshape(xx.shape)

    print("w1:\n",sess.run(w1))
    print("b1:\n",sess.run(b1))
    print("w2:\n",sess.run(w2))
    print("b2:\n",sess.run(b2))


plt.scatter(X[:,0], X[:,1], c=np.squeeze(Y_c))
plt.contour(xx,yy,probs,levels=[0.5])  # x坐标轴，y坐标轴，该点高度，等高线高度
plt.show()

# -----------------------------------------------------------------------
# 定义反向传播方法，含正则化
train_step = tf.train.AdamOptimizer(0.0001).minimize(loss_total)

with tf.Session() as sess:
    init_op = tf.global_variables_initializer()
    sess.run(init_op)
    STEPS = 40000
    for i in range(STEPS):
        start = (i*BATCH_SIZE) % 300
        end = start + BATCH_SIZE
        sess.run(train_step, feed_dict = {x:X[start:end], y_:Y_[start:end]})
        if i% 2000 == 0:
            loss_v = sess.run(loss_total, feed_dict ={x:X,y_:Y_})
            print ("在 %d 步之后，损失达到 %f" % (i, loss_v))

    xx, yy = np.mgrid[-3:3:.01, -3:3:.01]
    # 将xx，yy拉直，并合并成一个2列的矩阵，得到网格坐标点集合
    grid = np.c_[xx.ravel(),yy.ravel()]
    probs = sess.run(y, feed_dict = {x:grid})
    probs = probs.reshape(xx.shape)

    print("w1:\n",sess.run(w1))
    print("b1:\n",sess.run(b1))
    print("w2:\n",sess.run(w2))
    print("b2:\n",sess.run(b2))


plt.scatter(X[:,0], X[:,1], c=np.squeeze(Y_c))
plt.contour(xx,yy,probs,levels=[0.5])
plt.show()